// weakly informative priors on delta, D_0, sigma
// rhos unconstrained
// tail(rho, t-1) ~ normal(head(rho, t-1), 0.5) ;

functions {
 
   // Compute mean partisanship without error correction term
   
    real mean_wo_correction(real D_0, real rho, vector beta, 
                          real party_L1, real util, real econ) {
    real val ;
    val = D_0 + rho * party_L1 + 
      (1 - rho) * (beta[1] + beta[2] * util + beta[3] * econ) ;
    return val ;
  }
  
  // Compute error correction term
  real error_correction(real D_0, real rho_t1, vector beta, 
                        real party_L1, real party_L2, 
                        real util_L1, real econ_L1) {
    real val ; 
    val = party_L1 - D_0 - 
      rho_t1 * party_L2 - 
      (1-rho_t1) * (beta[1] + beta[2]*util_L1 + beta[3]*econ_L1) ;
    return val ;
  }
  
  // Compute mean partisanship with error correction
   
  real mean_w_correction(real delta, real D_0, real rho, real rho_t1, vector beta,
                         real party_L1, real party_L2, real util, real util_L1,
                         real econ, real econ_L1) {
    
    real m ; 
    real V ;
    m = mean_wo_correction(D_0, rho, beta, party_L1, util, econ) ;
    V = error_correction(D_0, rho_t1, beta, party_L1, party_L2, util_L1, econ_L1) ;
    return m - delta*V ;
  }
  
  
  // Transform standard normal to cauchy(location, scale).
   
   // @param location Location parameter for Cauchy distribution
   // @param scale Scale parameter for Cauchy distribution
   // @param noise Standard normal variable (as declared in parameters block)
   
  real cauchy_trans_lp(real location, real scale, real noise) {
    noise ~ normal(0,1) ;
    return location + scale * tan(pi() * (Phi_approx(noise) - 0.5)) ;
  }

  
   // Transform standard normal to normal(location, scale).
   
   // @param location Location parameter for normal distribution
   // @param scale Scale parameter for normal distribution
   // @param noise Standard normal variable (as declared in parameters block)
  
  real normal_trans_lp(real location, real scale, real noise) {
    noise ~ normal(0,1) ;
    return location + scale * noise ;
  }
  
  
   // Transform vector of standard normals to vector of normals
   // with same location (mean) and scale (sd).
   
   // @param location Location for normal distributions
   // @param scale Scale for normal distributions
   // @param noise Name of vector of standard normals (as declared in parameters block)
   
  vector normal_trans_vec_lp(real location, real scale, vector noise) {
    noise ~ normal(0,1) ;
    return location + scale * noise ;
  }

   // Compute inverse logit of vector elements
  
  vector inv_logit_vec(vector x) {
    vector[num_elements(x)] out ;
    for (j in 1:num_elements(out)) out[j] = inv_logit(x[j]) ;
    return out ;
  }
}
data {
  int<lower=1>    T ;         // number of time periods
  int<lower=1>    J ;         // number of groups
  matrix[T,J]     party ;     // mean partisanship  
  matrix[T,J]     party_L1 ;  // party, 1 lag
  matrix[T,J]     party_L2 ;  // party, 2 lags
  matrix[T,J]     util ;      // mean utility diff b/t parties on civil rights issues
  matrix[T,J]     util_L1 ;   // util, 1 lag
  vector[T]       econ ;      // retrospective economic evaluation
  vector[T]       econ_L1 ;   // econ, 1 lag
}
parameters {  
  vector[T]                 rho_logit;    // time-varying autoregressive parameter
  real                      D_0 ;         // constant
  real<lower=0>             sigma_noise ; // see below
  vector[3]                 beta_noise ;  // see below
  real                      delta_noise ; // see below
}
transformed parameters {
  vector<lower=0,upper=1>[T]  rho;      // time-varying autoregressive parameter
  real<lower=0>               sigma ;   // sd(party[t,j])
  vector[3]                   beta ;    // intercept & coefs on util, econ
  real                        delta ;   // coef on error correction term
  
  sigma = cauchy_trans_lp(0, 2.5, sigma_noise) ;
  delta = normal_trans_lp(0, 2.5, delta_noise) ;
  beta = normal_trans_vec_lp(0, 5, beta_noise) ;
  rho = inv_logit_vec(rho_logit) ;
  
}
model {
  // Priors
  D_0 ~ normal(0, 1) ;

  // Implied priors from transformed parameters block: 
  // beta[i] ~ normal(0, 5) ;
  // delta ~ normal(0, 2.5) ;
  // sigma ~ cauchy(0, 2.5) ;
  
  rho_logit[1] ~ normal(0,2.5) ;
  tail(rho_logit, T-1) ~ normal(head(rho_logit, T-1), 2.5) ;
  
  // Likelihood
  for (j in 1:J) { // loop over groups
    vector[T] mu ;
    mu[1] = mean_wo_correction(D_0, rho[1], beta, party_L1[1,j], util[1,j], econ[1]) ;
    for (t in 2:T) mu[t] = mean_w_correction(delta, D_0, rho[t], rho[t-1], beta,
                                             party_L1[t,j], party_L2[t,j], 
                                             util[t,j], util_L1[t,j],
                                             econ[t], econ_L1[t]) ;
    
    col(party, j) ~ normal(mu, sigma) ;
  }
}
generated quantities {
  matrix[T,J] mu_posterior ;    // posterior mean (fitted values)
  matrix[T,J] resids_fitted ;   // residuals: party[t,j] - mu_posterior[t,j]
  matrix[T,J] party_rep ;       // simulations from posterior predictive dist
  matrix[T,J] resids_rep ;      // residuals: party[t,j] - party_rep[t,j]
  vector[T*J] log_lik ;         // log likelihood (for computing waic)

  { // local block 
  matrix[T,J] log_lik_mat ;
  for (j in 1:J) {
    for(t in 1:T) {
      if (t == 1) 
        mu_posterior[t,j] = mean_wo_correction(D_0, rho[t], beta, 
                                party_L1[t,j], util[t,j], econ[t]) ;
      else 
        mu_posterior[t,j] = mean_w_correction(delta, D_0, 
                                rho[t], rho[t-1], beta,
                                party_L1[t,j], party_L2[t,j], 
                                util[t,j], util_L1[t,j],
                                econ[t], econ_L1[t]) ;
      
      party_rep[t,j] = normal_rng(mu_posterior[t,j], sigma) ;      
      log_lik_mat[t,j] = normal_log(party[t,j], mu_posterior[t,j], sigma) ;
    }
  }
  
  log_lik = to_vector(log_lik_mat) ;
  } //end local block

  resids_rep = party - party_rep ; 
  resids_fitted = party - mu_posterior ;
}
